from environment.uav import UAV
from environment.target import Target
from environment.multi_target import MultiTarget
import numpy as np
import math
import random

class HierarchicalEncirclementEnv:
    """分层决策的多目标围捕环境"""
    
    def __init__(self, num_uavs, grid_size, num_position_action, num_energy_action, num_targets=3, seed=None, 
                 observation_range=150.0, communication_range=200.0, partial_observability=True):
        if seed is not None:
            np.random.seed(seed)
            random.seed(seed)
            
        self.grid_size = grid_size
        self.num_uavs = num_uavs
        self.num_targets = num_targets
        
        # 部分可观测性参数
        self.observation_range = observation_range
        self.communication_range = communication_range
        self.partial_observability = partial_observability
        
        # 创建无人机
        self.uavs = []
        for i in range(num_uavs):
            uav = UAV(grid_size, i, [], num_position_action, num_energy_action)
            # 分散初始化位置
            angle = 2 * np.pi * i / num_uavs
            radius = grid_size // 3
            center_x, center_y = grid_size // 2, grid_size // 2
            uav.position = (
                center_x + radius * np.cos(angle),
                center_y + radius * np.sin(angle)
            )
            self.uavs.append(uav)
        
        # 创建多个动态目标
        self.multi_target = MultiTarget(grid_size, num_targets=num_targets)
        
        # 目标分配跟踪
        self.target_assignments = {}  # {uav_id: target_id}
        self.target_assignment_history = []
        
        # 围捕状态跟踪
        self.target_encirclement_status = [False] * num_targets
        self.encircled_targets = set()
        
        # 奖励参数
        self.ASSIGNMENT_REWARD = 10
        self.ENCIRCLEMENT_REWARD = 200
        self.COOPERATION_REWARD = 50
        self.DISTANCE_REWARD_SCALE = 30
        self.COLLISION_PENALTY = -20
    
    def get_observable_entities(self, observer_uav_id):
        """获取指定无人机可观测的实体"""
        if not self.partial_observability:
            # 如果不启用部分可观测性，返回全部信息
            return {
                'uavs': list(range(self.num_uavs)),
                'targets': list(range(self.num_targets)),
                'communications': list(range(self.num_uavs))
            }
        
        observer_pos = np.array(self.uavs[observer_uav_id].position)
        observable_entities = {
            'uavs': [observer_uav_id],  # 总是能观测到自己
            'targets': [],
            'communications': [observer_uav_id]  # 总是能与自己通信
        }
        
        # 检查可观测的其他无人机
        for uav_id, uav in enumerate(self.uavs):
            if uav_id == observer_uav_id:
                continue
                
            uav_pos = np.array(uav.position)
            distance = np.linalg.norm(observer_pos - uav_pos)
            
            # 观测范围内的无人机
            if distance <= self.observation_range:
                observable_entities['uavs'].append(uav_id)
            
            # 通信范围内的无人机
            if distance <= self.communication_range:
                observable_entities['communications'].append(uav_id)
        
        # 检查可观测的目标
        for target_id, target in enumerate(self.multi_target.targets):
            target_pos = np.array(target['position'])
            distance = np.linalg.norm(observer_pos - target_pos)
            
            if distance <= self.observation_range:
                observable_entities['targets'].append(target_id)
        
        return observable_entities
    
    def get_partial_state(self, observer_uav_id):
        """获取指定无人机的部分观测状态"""
        observable = self.get_observable_entities(observer_uav_id)
        state = []
        
        # 可观测的无人机状态
        for uav_id in observable['uavs']:
            uav = self.uavs[uav_id]
            state.extend([uav.position[0], uav.position[1], uav.energy])
        
        # 不可观测的无人机用占位符填充
        unobservable_uavs = self.num_uavs - len(observable['uavs'])
        for _ in range(unobservable_uavs):
            state.extend([-1, -1, -1])  # 使用-1表示未知信息
        
        # 可观测的目标状态
        for target_id in observable['targets']:
            target = self.multi_target.targets[target_id]
            state.extend([target['position'][0], target['position'][1]])
        
        # 不可观测的目标用占位符填充
        unobservable_targets = self.num_targets - len(observable['targets'])
        for _ in range(unobservable_targets):
            state.extend([-1, -1])  # 使用-1表示未知信息
        
        # 通信范围内的分配状态
        assignment_state = [0] * (self.num_uavs * self.num_targets)
        for uav_id in observable['communications']:
            if uav_id in self.target_assignments:
                target_id = self.target_assignments[uav_id]
                if target_id < self.num_targets:
                    idx = uav_id * self.num_targets + target_id
                    assignment_state[idx] = 1
        
        state.extend(assignment_state)
        
        # 添加可观测性信息（用于调试和分析）
        observability_info = [
            len(observable['uavs']) / self.num_uavs,  # 可观测无人机比例
            len(observable['targets']) / self.num_targets,  # 可观测目标比例
            len(observable['communications']) / self.num_uavs  # 通信连接比例
        ]
        state.extend(observability_info)
        
        return np.array(state)
    
    def reset(self, seed=None):
        """重置环境"""
        if seed is not None:
            np.random.seed(seed)
            random.seed(seed)
        
        # 重置目标
        self.multi_target.reset()
        
        # 重置无人机
        for i, uav in enumerate(self.uavs):
            angle = 2 * np.pi * i / self.num_uavs
            radius = self.grid_size // 3
            center_x, center_y = self.grid_size // 2, self.grid_size // 2
            uav.position = (
                center_x + radius * np.cos(angle),
                center_y + radius * np.sin(angle)
            )
            uav.history_position = []
            uav.energy = 100
        
        # 重置分配状态
        self.target_assignments.clear()
        self.target_assignment_history.clear()
        self.target_encirclement_status = [False] * self.num_targets
        self.encircled_targets.clear()
    
    def update_target_assignment(self, uav_id, target_id):
        """更新目标分配"""
        self.target_assignments[uav_id] = target_id
        self.target_assignment_history.append((uav_id, target_id))
    
    def check_target_encirclement(self, target_id):
        """检查指定目标是否被围捕"""
        if target_id >= len(self.multi_target.targets):
            return False
        
        target_pos = np.array(self.multi_target.targets[target_id]['position'])
        
        # 统计围绕该目标的无人机数量
        surrounding_uavs = 0
        for uav_id, assigned_target in self.target_assignments.items():
            if assigned_target == target_id:
                uav_pos = np.array(self.uavs[uav_id].position)
                distance = np.linalg.norm(uav_pos - target_pos)
                if distance <= 50:  # 围捕距离阈值
                    surrounding_uavs += 1
        
        # 需要至少3个无人机围捕一个目标
        return surrounding_uavs >= 3
    
    def check_all_encirclements(self):
        """检查所有目标的围捕状态"""
        newly_encircled = []
        
        for target_id in range(self.num_targets):
            is_encircled = self.check_target_encirclement(target_id)
            
            if is_encircled and target_id not in self.encircled_targets:
                newly_encircled.append(target_id)
                self.encircled_targets.add(target_id)
                self.target_encirclement_status[target_id] = True
        
        return newly_encircled
    
    def calculate_hierarchical_reward(self, uav_id, assigned_target_id, current_step, max_steps):
        """计算分层决策的奖励"""
        uav = self.uavs[uav_id]
        uav_pos = np.array(uav.position)
        
        total_reward = 0
        
        # 1. 目标分配奖励
        if assigned_target_id < len(self.multi_target.targets):
            target_pos = np.array(self.multi_target.targets[assigned_target_id]['position'])
            distance_to_assigned_target = np.linalg.norm(uav_pos - target_pos)
            
            # 距离奖励（越近越好）
            max_distance = np.sqrt(2) * self.grid_size
            distance_reward = self.DISTANCE_REWARD_SCALE * (1 - distance_to_assigned_target / max_distance)
            total_reward += distance_reward
        
        # 2. 围捕成功奖励
        newly_encircled = self.check_all_encirclements()
        if assigned_target_id in newly_encircled:
            total_reward += self.ENCIRCLEMENT_REWARD
        
        # 3. 协作奖励
        cooperation_reward = 0
        for other_uav_id, other_target_id in self.target_assignments.items():
            if other_uav_id != uav_id and other_target_id == assigned_target_id:
                cooperation_reward += self.COOPERATION_REWARD / self.num_uavs
        total_reward += cooperation_reward
        
        # 4. 避碰惩罚
        collision_penalty = 0
        for other_uav in self.uavs:
            if other_uav != uav:
                other_pos = np.array(other_uav.position)
                if np.linalg.norm(uav_pos - other_pos) < 20:
                    collision_penalty += self.COLLISION_PENALTY
        total_reward += collision_penalty
        
        # 5. 时间效率奖励
        if len(newly_encircled) > 0:
            efficiency_ratio = (max_steps - current_step) / max_steps
            total_reward += 50 * efficiency_ratio
        
        return total_reward
    
    def perform_hierarchical_action(self, uav_id, action, assigned_target_id, current_step=0, max_steps=150):
        """执行分层动作并返回奖励"""
        uav = self.uavs[uav_id]
        
        # 执行无人机移动
        movement_reward, _ = uav.execute_cruise_action(action)
        
        # 移动所有目标
        self.multi_target.move_all_targets()
        
        # 更新目标分配
        self.update_target_assignment(uav_id, assigned_target_id)
        
        # 计算分层奖励
        hierarchical_reward = self.calculate_hierarchical_reward(
            uav_id, assigned_target_id, current_step, max_steps
        )
        
        total_reward = movement_reward + hierarchical_reward
        return total_reward
    
    def get_global_state(self):
        """获取全局状态"""
        state = []
        
        # 无人机状态
        for uav in self.uavs:
            state.extend([uav.position[0], uav.position[1], uav.energy])
        
        # 目标状态
        for target in self.multi_target.targets:
            state.extend([target['position'][0], target['position'][1]])
        
        # 分配状态
        assignment_state = [0] * (self.num_uavs * self.num_targets)
        for uav_id, target_id in self.target_assignments.items():
            if target_id < self.num_targets:
                idx = uav_id * self.num_targets + target_id
                assignment_state[idx] = 1
        state.extend(assignment_state)
        
        return np.array(state)
    
    def is_done(self):
        """检查任务是否完成"""
        return len(self.encircled_targets) >= self.num_targets
    
    def get_success_rate(self):
        """获取成功率"""
        return len(self.encircled_targets) / self.num_targets